#include <stdio.h>
#include <string.h>

#include "viterbi.h"

#define MAX_STAT_NUM	16

/* evaluation */
//finding the probability of an observed sequence given a HMM
double forward(unsigned int * seq, int n)
{
	if(n <= 0 || seq == NULL){
		return 0;
	}
	double trellis[HNUM][MAX_STAT_NUM];
	memset(trellis, 0, sizeof(double) * HNUM * n);

	/* calc partial probility */
	double sum = 0;
	int i;
	int j;
	for(j = 0; j < n + 1; j ++){
		for(i = SUNNY; i <= RAINY; i ++){
			if(j == 0){
				trellis[0][i] = INIT[i] * CONF[i][seq[0]];
				printf("trellis[0][%d] = %f * %f = %f\n",
					i, INIT[i], CONF[i][seq[0]], trellis[0][i]);
			}else if(j == n){
				sum += trellis[n-1][i];
			}else if(j > 0){
				int k;
				double sumk = 0;
				for(k = SUNNY; k <= RAINY; k ++){
					sumk += trellis[j-1][k] * TRANS[k][i];
					printf("%u\t+ %f * %f\n",
							k, trellis[j-1][k], TRANS[k][i]);
				}
				trellis[j][i] = sumk * CONF[i][seq[j]];
				printf("trellis[%d][%d] = ( ) * %f = %f\n",
					j, i, CONF[i][seq[j]], trellis[j][i]);
			}
		}
	}

	return sum;
}

typedef struct vt{
	double p;
	unsigned int path[MAX_STAT_NUM];
}vt;
/* decoding */
//finding the most probable sequence of hidden states given some observations
int viterbi(unsigned int * seq, int n, unsigned int * code)
{
	if(n <= 0 || seq == NULL || code == NULL){
		return 0;
	}
	vt trellis[HNUM][MAX_STAT_NUM];
	memset(trellis, 0, sizeof(vt) * HNUM * n);

	/* calc partial probility */
	double max = 0;
	int i;
	int j;
	for(j = 0; j < n + 1; j ++){
		for(i = SUNNY; i <= RAINY; i ++){
			if(j == 0){
				trellis[0][i].p = INIT[i] * CONF[i][seq[0]];
				printf("trellis[0][%d] = %f * %f = %f\n",
					i, INIT[i], CONF[i][seq[0]], trellis[0][i].p);
			}else if(j == n){
				if(max < trellis[j-1][i].p){
					max = trellis[j-1][i].p;
					memcpy(code, trellis[j-1][i].path, 
							(j - 1) * sizeof(unsigned int));
					code[j-1] = i;
				}
			}else if(j > 0){
				int k;
				double maxk = 0;
				for(k = SUNNY; k <= RAINY; k ++){
					if(maxk < trellis[j-1][k].p * TRANS[k][i]){
						maxk = trellis[j-1][k].p * TRANS[k][i];
						memcpy(trellis[j][i].path, trellis[j-1][i].path, 
								(j - 1) * sizeof(unsigned int));
						trellis[j][i].path[j-1] = k;
					}
					printf("%u\t~ %f * %f = %f\n",
							k, trellis[j-1][k].p, TRANS[k][i], 
							trellis[j-1][k].p * TRANS[k][i]);
				}
				trellis[j][i].p = maxk * CONF[i][seq[j]];
				printf("trellis[%d][%d] = ( ) * %f = %f, best %u\n",
					j, i, CONF[i][seq[j]], trellis[j][i].p,
					trellis[j][i].path[j-1]);
			}
		}
	}

	return 0;
}

int main()
{
	unsigned int seq[] = {DRY, DAMP, SOGGY};

	double p = forward(seq, sizeof(seq)/sizeof(unsigned int));
	printf("forward %f\n", p);

	unsigned int code[sizeof(seq)/sizeof(unsigned int)];
	viterbi(seq, sizeof(seq)/sizeof(unsigned int), code);
	int i;
	for(i = 0; i < sizeof(seq)/sizeof(unsigned int); i ++){
		printf("%u\t", code[i]);
	}
	printf("\n");
	return 0;
}

